# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
import time
from typing import Literal
from tortoise.expressions import Q
from common.utils.times import TimeUtil
from common.utils.config import ConfigUtil
from common.models.users import UserModel
from common.models.agents import AgentMcpModel
from common.models.agents import AgentKnowModel
from common.models.agents import AgentRobotModel
from common.models.draws import DrawsRecordsModel
from common.models.video import VideoRecordsModel
from common.models.music import MusicRecordsModel
from common.models.users import UserSubscribeModel
from common.models.market import SubscribePackageModel
from common.models.market import SubscribePrivilegeModel
from common.postgres.public import PgKnowledgeModel


class SubscribeLogic:

    @classmethod
    async def get_user_vip(cls, user_id: int) -> dict:
        """
        获取用户最高VIP

        Args:
            user_id: int  用户ID

        Returns:
            dict

        Author:
            zero
        """
        # 按生效时间查询会员 (未过期会员)
        members = await (UserSubscribeModel
                         .filter(
                                Q(
                                    user_id=user_id,
                                    is_perpetual=1,
                                    is_delete=0
                                ) |
                                Q(
                                    user_id=user_id,
                                    is_perpetual=0,
                                    is_delete=0,
                                    expire_time__gt=int(time.time())
                                )
                         ).order_by("-effect_time").all().values())

        # 会员是否过期
        is_member_overdue = False

        # 无会员的记录
        if not members:
            basis = await SubscribePackageModel.filter(is_builtin=1, is_delete=0).first()
            if not basis:
                return {}
            members = [{"id": 0, "package_id": basis.id, "is_perpetual": 1, "effect_time": 0, "expire_time": 0}]
        else:
            current_time = int(time.time())
            remain_time = members[0]["remain_time"]
            effect_time = members[0]["effect_time"]
            if remain_time > 0:
                remain_time = 0
            if effect_time > current_time:
                effect_time = current_time
            if remain_time == 0 or members[0]["effect_time"] > current_time:
                await UserSubscribeModel.filter(id=members[0]["id"]).update(
                    remain_time=remain_time,
                    effect_time=effect_time,
                    update_time=current_time
                )

        # 返回生效会员
        package = await SubscribePackageModel.filter(id=members[0]["package_id"]).first()
        return {
            "id": package.id,
            "name": package.name,
            "description": package.description,
            "train_num": package.train_num,
            "robot_num": package.robot_num,
            "know_num": package.know_num,
            "member_id": members[0]["id"],
            "end_time": "永久" if members[0]["is_perpetual"] else TimeUtil.timestamp_to_date(members[0]["expire_time"]),
            "effect_time": members[0]["effect_time"],
            "expire_time": members[0]["expire_time"],
            "is_perpetual": members[0]["is_perpetual"],
            "is_builtin": True if package.is_builtin else False,
            "is_overdue": is_member_overdue
        }

    @classmethod
    async def get_user_quota(cls, user_id: int, field: Literal["robot", "know", "train", None] = None) -> dict:
        """
        获取用户配额概况

        Args:
            user_id: int (用户ID)
            field: Literal["robot", "know", "train", None] (查询字段)

        Returns:
            Union[Dict[int], Dict[str, Dict[int]]]

        Author:
            zero
        """
        # 查询当前订阅
        member = await cls.get_user_vip(user_id)

        # 获取信息
        user = await UserModel.filter(id=user_id).first()
        pack = await SubscribePackageModel.filter(id=int(member["id"])).first()

        # 总共合计配额
        total_know_num: int = pack.know_num + user.know_num
        total_robot_num: int = pack.robot_num + user.robot_num
        total_train_num: int = pack.train_num + user.train_num
        total_mcp_num: int = pack.mcp_num + user.mcp_num

        # 定义数据结构
        details = {
            "mcp": {
                "pack_num": pack.mcp_num,
                "user_num": user.mcp_num,
                "total_num": total_mcp_num,
                "usage_num": 0,
                "surplus_num": 0,
            },
            "know": {
                "pack_num": pack.know_num,
                "user_num": user.know_num,
                "total_num": total_know_num,
                "usage_num": 0,
                "surplus_num": 0,
            },
            "robot": {
                "pack_num": pack.robot_num,
                "user_num": user.robot_num,
                "total_num": total_robot_num,
                "usage_num": 0,
                "surplus_num": 0
            },
            "train": {
                "pack_num": pack.train_num,
                "user_num": user.train_num,
                "total_num": total_train_num,
                "usage_num": 0,
                "surplus_num": 0
            },
            "member": member
        }

        # 知识库使用数量
        if field is None or field == "mcp":
            mcp_use_num: int = await AgentMcpModel.filter(user_id=user_id, is_delete=0).count()
            details["mcp"]["usage_num"] = mcp_use_num
            details["mcp"]["surplus_num"] = total_mcp_num - mcp_use_num
        # 知识库使用数量
        if field is None or field == "know":
            know_use_num: int = await AgentKnowModel.filter(user_id=user_id, is_delete=0).count()
            details["know"]["usage_num"] = know_use_num
            details["know"]["surplus_num"] = total_know_num - know_use_num
        # 智能体使用数量
        if field is None or field == "robot":
            robot_use_num: int = await AgentRobotModel.filter(user_id=user_id, is_delete=0).count()
            details["robot"]["usage_num"] = robot_use_num
            details["robot"]["surplus_num"] = total_know_num - robot_use_num
        # 训练组使用数量
        if field is None or field == "train":
            train_use_num: int = await PgKnowledgeModel.filter(user_id=user_id, is_delete=0).count()
            details["train"]["usage_num"] = train_use_num
            details["train"]["surplus_num"] = total_train_num - train_use_num

        # 返回最终的结果
        return details if field is None else details[field]

    @classmethod
    async def get_surplus_quota(cls, user_id: int, scene: Literal["robot", "know", "train", "mcp"]) -> int:
        """
        获取套餐剩余配额

        Args:

            user_id: int (用户ID)
            scene: Literal["robot", "know", "train"] (查询场景)

        Returns:
            int: 剩余的数量

        Author:
            zero
        """
        # 查询当前订阅
        member = await cls.get_user_vip(user_id)
        if not member:
            return -1

        # 获取套餐信息
        package = await SubscribePackageModel.filter(id=int(member["id"])).first()

        # 获取用户信息
        user = await UserModel.filter(id=user_id).first()

        # 总共合计配额
        total_know_num: int = package.know_num + user.know_num
        total_robot_num: int = package.robot_num + user.robot_num
        total_train_num: int = package.train_num + user.train_num
        total_mcp_num: int = package.mcp_num + user.mcp_num

        if scene == "mcp":
            use_num: int = await AgentMcpModel.filter(user_id=user_id, is_delete=0).count()
            return total_mcp_num - use_num
        elif scene == "know":
            use_num: int = await AgentKnowModel.filter(user_id=user_id, is_delete=0).count()
            return total_know_num - use_num
        elif scene == "robot":
            use_num: int = await AgentRobotModel.filter(user_id=user_id, is_delete=0).count()
            return total_robot_num - use_num
        elif scene == "train":
            use_num: int = await PgKnowledgeModel.filter(user_id=user_id, is_delete=0).count()
            return total_train_num - use_num
        else:
            return 0

    @classmethod
    async def get_surplus_free(cls, user_id: int, scene: str, mark: str = None):
        """
        获取套餐剩余免费次数
        """
        member = await cls.get_user_vip(user_id)
        if not member:
            return []

        package_id: int = int(member["id"])

        # [AI绘画/AI视频/AI音乐]
        if scene in ["draws", "video", "music"] and mark:
            _scene = scene + "_" + mark
            privilege = await SubscribePrivilegeModel.filter(scene=_scene, package_id=package_id).first()
            if not privilege:
                from apps.admin.service.market.subscribe_service import SubscribeService
                await SubscribeService.detail(package_id)
                privilege = await SubscribePrivilegeModel.filter(scene=_scene, package_id=package_id).first()

            model_cls = {"draws": DrawsRecordsModel, "video": VideoRecordsModel, "music": MusicRecordsModel}
            usage_num: int = await (model_cls[scene]
                                    .filter(user_id=user_id)
                                    .filter(engine=mark)
                                    .filter(is_free=1)
                                    .filter(status__in=[0, 1, 2])  # 等待,运行,成功
                                    .filter(create_time__gte=TimeUtil.today()[0])
                                    .filter(create_time__lte=TimeUtil.today()[1])
                                    .count())

            free_status: bool = False
            surplus_num: int = int(privilege.daily_limit - usage_num)
            if not privilege.is_disable and (privilege.daily_limit <= 0 or surplus_num > 0):
                free_status = True

            configs = await ConfigUtil.get(scene) or {}
            name = configs.get(mark, {}).get("name")
            if not configs.get("name"):
                opt = {"draws": "绘画", "video": "视频", "music": "音乐"}
                name = mark.upper() + opt[scene]

            return {
                "type": scene,
                "mark": mark,  # 应用标识: {mj,sd,dalle,doubao}, {luma,zhipu}, {suno}
                "name": name,  # 应用名称
                "usage_num": usage_num,  # 当日已用免费次数
                "surplus_num": surplus_num,  # 当日剩余免费次数
                "daily_limit": privilege.daily_limit,  # 每日限制免费次数: [0=免费且不限制]
                "free_status": free_status,  # 当前是否免费状态: [True=免费, False=不免费]
                "is_disable": privilege.is_disable  # 是否禁止使用: [1=是, 0=否]
            }
        return None
